/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/******************************************************************************************************
 * stepMain.c - functions and function-calls related to various operations performed at every MD step *
 ******************************************************************************************************/

// #define DEBUG // uncomment this line to turn on debugging functionalities

#include "dataStructs.h"

extern VecR *siteCoords, *siteVelocity, *siteForce, *siteTorque, *atomCoordinates, *currWatCoords, *prevWatCoords;
extern Site *site;  
extern Atom *atom;  
extern const LipidStruct dopcStruct, dopeStruct;
extern VecR region, vSum;  
extern const real *mass, deltaT, cholCharge, totMassExceptConstrSolute, totalMass; 
extern real timeNow;
extern Prop totEnergy, potEnergy, zConstraintForce;
extern const int applyThermostat, applyBarostat, flexBox, removeSystemTranslation, removeMonolayersTranslation;
extern const int nSites, nWaters, nLipids, nAtoms, nSolutes, DOFs, zConstraint, stepNebr;
extern const int stepPdb, reCenterBilayer, stepLimit;
extern int stepCount; 
extern RMat boxMatrix;

// Global variables:
VecRTwo velSumUpperLayer, velSumLowerLayer; // net velocity vectors of each lipid layer, 
                                            // used in lipid lat diff measurements, initialised in startUp.c
const int stepAvg; /* number of timesteps between evaluations of averages (controls averaging "frequency") */ 

/* External function declarations: */
void AccumProps( int );
void ApplyBarostat( RMat );
void ApplyThermostat();
void BuildNeighborList();
void CheckPressureTensor( RMat );
void ComputeAngleBendInteractions( real*, RMat* );
void ComputeNonBondedInteractions( real*, RMat* );
void ComputeBondStretchInteractions( real*, RMat* );
void ComputeOrientationRestraints( real* );
void CountPossiblePermeatingWaters();
void CountWatersPerSection();
void DoNeighborList();
void EvalProps( RMat, RMat* );
void InvertMatrix( RMat*, RMat* );
void LeapfrogStep( int );
void NvtStep( int );
void NptStep( int );
void PrintAvgZConstraintParameters();
void PrintMatrix( RMat );
void PrintSplitTemperatures( FILE* );
void PrintSummary( FILE* );
void PrintSummaryOnFile();
void PrintSummaryWater( FILE* );
void PrintSummaryWaterOnFile();
void PrintTailOrderParametersOnFile();
void PrintWatLatMotion();
void PrintZConstraintForce();
void PutCheckpoint();
void UpdateCellSize();

static void ApplyBoundaryCond() // takes care of coordinate wraparound [Rapaport, p.27] */
{
  int n; 
  DO_SITE VWrapAll( siteCoords[ n ] ); /* macro VWrapAll is defined in "definitions.h" */
}

static void ApplyBoundaryCondS() 
{
  int n; 
  DO_SITE VWrapAllS( siteCoords[ n ] );
}

/**********************************************************************************************************************************
 * RemoveMonolayersTranslation() -- for each of the two monolayers, removes the mass center motion in the xy plane (the bilayer 
 * plane by convention). This prevents the monolayers from "drifting" laterally, which would be unrealistic and would introduce 
 * artefacts in the calculation of the lipid diffusion. Drift-removal is suggested by Patra et al [J Phys Chem B, 108, 4485 (2004)]. 
 * However, Klauda et al. argue against it, and show that for a large enough bilayer (288 lipids) the effect is negligible [J Chem 
 * Phys, 125, 144710 (2006)]. Falck et al. and Lindahl & Edholm did not remove the monolayers' translation but calculated the 
 * lateral diffusion coefficients relative to the monolayers' mass centers, thus correcting for any drifts [Biophys. J., 87, 1076 
 * (2004); J Chem Phys, 115, 4938 (2001)]. Roark & Feller calculated diffusion coefficients both with and without removal of mono-
 * layers translation, showing size dependency for small systems (18-72 lipids) but similar results for large bilayers (288 lipids)
 * [J Phys Chem B, 113, 13229 (2009)].
 **********************************************************************************************************************************/
static void RemoveMonolayersTranslation() 
{ 
  int n, cholIndex, sCount, nLipidsUpperMonolayer = 0, nLipidsLowerMonolayer = 0;
  real scalingFactorUpper, scalingFactorLower;
  VecRTwo comVelUpperLayer;
  VecRTwo comVelLowerLayer;

  VZeroTwo( comVelUpperLayer ); // centre of mass velocity of the upper monolayer
  VZeroTwo( comVelLowerLayer ); // centre of mass velocity of the lower monolayer
  
  //  computing xy-COM vels
  DO_SITE { // scan all sites
    switch ( site[ n ].lipidSpecies ) {
    case DOPC:
    case DOPE: // same as DOPC because they have the same number of sites (see dataStructs.h)
      cholIndex = n; /* index of choline unit for current lipid (assumes topology defined in dataStructs.h) */ 
      if ( siteCoords[ cholIndex + 2 ].z > 0 ) { // if glycerol site of current lipid is in the "upper" layer 
	nLipidsUpperMonolayer++;
	for ( sCount = cholIndex; sCount < ( cholIndex + dopcStruct.nBeads ); sCount++ ) // scan through all sites belonging to the current lipid
	  VVAddTwo( comVelUpperLayer, siteVelocity[ sCount ] );  // accumulate velocity
      } else { // "lower" layer
	nLipidsLowerMonolayer++;
	for ( sCount = cholIndex; sCount < ( cholIndex + dopcStruct.nBeads ); sCount++ ) // scan through all sites belonging to the current lipid
	  VVAddTwo( comVelLowerLayer, siteVelocity[ sCount ] );
      }
      n += ( dopcStruct.nBeads - 1 ); // update index - "-1" is necessary as DO_SITE already causes a unitary increment
      break;
    case NON_LIPID: break;
    default: printf("ERROR in RemoveMonolayersTranslation()! Exiting...\n"); exit(0);
    }
  }

  scalingFactorUpper = - 1. / ( nLipidsUpperMonolayer * dopcStruct.nBeads );
  scalingFactorLower = - 1. / ( nLipidsLowerMonolayer * dopcStruct.nBeads );

  // removing net lateral vel
  DO_SITE {
    switch ( site[ n ].lipidSpecies ) {
    case DOPC:
    case DOPE:
      cholIndex = n; 
      if ( siteCoords[ cholIndex + 2 ].z > 0 ) {
	for ( sCount = cholIndex; sCount < ( cholIndex + dopcStruct.nBeads ); sCount++ )
	  VVSAddTwo( siteVelocity[ sCount ], scalingFactorUpper, comVelUpperLayer );
      } else {
	for ( sCount = cholIndex; sCount < ( cholIndex + dopcStruct.nBeads ); sCount++ )
	  VVSAddTwo( siteVelocity[ sCount ], scalingFactorLower, comVelLowerLayer );
      }
      n += ( dopcStruct.nBeads - 1 ); // update index - "-1" is necessary as DO_SITE already causes a unitary increment
      break;
    case NON_LIPID: break;
    default: printf("ERROR in RemoveMonolayersTranslation()! Exiting...\n"); exit(0);
    }
  }
  // CHECKING

  /*   VZeroTwo( comVelUpperLayer ); */
  /*   VZeroTwo( comVelLowerLayer ); */
  
  /*   // checking proper removal */
  /*   DO_LIPID {  */
  /*     cholIndex = n * N_SITES_PER_DMPC;  */
  /*     if ( siteCoords[ cholIndex ].z > 0 ) {// upper layer, considering where the choline is */
  /*       for ( sCount = cholIndex; sCount < ( cholIndex + N_SITES_PER_DMPC ); sCount++ ) */
  /* 	VVAddTwo( comVelUpperLayer, siteVelocity[ sCount ] ); */
  /*     } else { */
  /*       for ( sCount = cholIndex; sCount < ( cholIndex + N_SITES_PER_DMPC ); sCount++ ) */
  /* 	VVAddTwo( comVelLowerLayer, siteVelocity[ sCount ] ); */
  /*     } */
  /*   } */

  /*   printf("*** After xy-com removal:\n"); */
  /*   printf("***  comVelUpperLayer = %4.3f %4.3f\n", comVelUpperLayer.x, comVelUpperLayer.y ); */
  /*   printf("***  comVelLowerLayer = %4.3f %4.3f\n\n", comVelLowerLayer.x, comVelLowerLayer.y ); */

}

/****************************************************************************************************
 * RemoveSystemTranslation() -- Remove overall net translational motion of the system mass centre
 *                              [Harvey et al, J Comp Chem, 19, 726 (1998); Gromacs user manual 
 *                               version 3.1.1, section 3.4] 
 ****************************************************************************************************/
static void RemoveSystemTranslation() 
{
  int n;
  int watersUpperSection = 0, watersLowerSection = 0;

  real sFac;

  VecR velSum;

  VZero( velSum );
  DO_SITE VVAdd( velSum, siteVelocity[ n ] ); 

  if ( reCenterBilayer ) { /* recentering the bilayer **this should go in a separate function** */
    DO_SITE {
      if ( WATER == site[ n ].lipidUnit ) { 
	if ( siteCoords[ n ].z > 0 ) watersUpperSection++;
	if ( siteCoords[ n ].z < 0 ) watersLowerSection++;
      }
    }
    
    if ( watersUpperSection >  watersLowerSection ) { // move water from top to bottom
      if ( velSum.z < 0 ) { // if the bilayer dritfs down
	DO_SITE VVSAdd( siteVelocity[ n ], - 1. / nSites, velSum ); // pull it back
      } // if it drifts up let it go, so waters get redistributed as we want
    } else { // if watersUpperSection <  watersLowerSection  move water from bottom to top
      if ( velSum.z > 0 ) { // if the bilayer dritfs up
	DO_SITE VVSAdd( siteVelocity[ n ], - 1. / nSites, velSum );
      }
    } // end else
    
    if ( abs( watersUpperSection - watersLowerSection ) == 1 ) {
      printf("Bilayer recentered: %d waters in the upper half of region and %d in the lower half\n", watersUpperSection, watersLowerSection++ );
      printf("Putting checkpoint...\n");
      PutCheckpoint ();
      printf("Exiting...\n");
      exit(0);
    }  
  } else { // if not reCenterBilayer - proper net translation removal
 
    // remove total net linear momentum
    sFac = - 1. / totalMass; // printf("totMass = %f amu\n", totalMass* MASS_amu);
    
    DO_SITE VVSAdd( velSum, mass[ site[ n ].type ], siteVelocity[ n ] ); // velSum is really the total linear momentum
    DO_SITE VVSAdd( siteVelocity[ n ], sFac, velSum );

    /*     if ( zConstraint ) { */
    /*       sFac = - 1. / ( nSites - 1 ); // assuming the solute is the last element of the site[] array */
    /*       for ( n = 0; n < nSites - 1; n++ ) VVAdd( velSum, siteVelocity[ n ] ); */
    /*       for ( n = 0; n < nSites - 1; n++ ) VVSAdd( siteVelocity[ n ], sFac, velSum ); */
    /*     } else { */
    
    /*       sFac = - 1. / nSites; */
    /*       DO_SITE VVAdd( velSum, siteVelocity[ n ] ); */
    /*       DO_SITE VVSAdd( siteVelocity[ n ], sFac, velSum ); */
    
    /*     } */
  }
}

static void ScaleCoords() 
{
  int n; 
  RMat invBoxMat;
  VecR s;

  InvertMatrix( &invBoxMat, &boxMatrix );
  
  DO_SITE {
    //    if ( flexBox ) { // anisotropic region 
    MVMul( s, invBoxMat.u, siteCoords[ n ] );
    VCopy( siteCoords[ n ], s );
    //} else VDiv( siteCoords[ n ], siteCoords[ n ], region ); 
    // printf(" site[%3d] = % f, % f, % f\n", n, siteCoords[ n ].x,siteCoords[ n ].y,siteCoords[ n ].z); 
  }
}

static void UnscaleCoords() 
{
  int n;
  VecR r;

  DO_SITE {
    //if ( flexBox ) {
    MVMul( r, boxMatrix.u, siteCoords[ n ] );
    VCopy( siteCoords[ n ], r );
    //} else VMul( siteCoords[ n ], siteCoords[ n ], region  );
  }
}

/********************************************************************************
 * StoreWatCoordsPrevStep() -- store water coordinates at the previous timestep *
 ********************************************************************************/
static void StoreWaterCoorsPrevStep() 
{ 
  int n, i=0;
  DO_SITE { 
    if ( WATER == site[ n ].lipidUnit ) { 
      VCopy( prevWatCoords [ i ], siteCoords[ n ] );
      i++;
    }
  }
}


/*******************************************************************************
 * StoreWatCoordsCurrStep() -- store water coordinates at the current timestep *
 *******************************************************************************/
static void StoreWatCoordsCurrStep() 
{
  int n, i=0;
  DO_SITE { 
    if ( WATER == site[ n ].lipidUnit ) { 
      VCopy( currWatCoords[ i ], siteCoords[ n ] );
      i++;
    }
  }
  CountPossiblePermeatingWaters();
}

/**************************************************************************************************************
 * SingleStep -- handles the processing for a single time step, including calls to functions dealing with     *
 * force evaluation, integration of the equations of motion, PBC, measurements. Adapted from [Rapaport, p.21] *
 **************************************************************************************************************/
void SingleStep()
{
  int i, n;
  real vzConst, fzConst;
  RMat systemPressure, systemVirial;

  ++stepCount;
  timeNow = stepCount * deltaT;

  MZero( systemPressure.u ); /* initialise pressure tensor for a new estimation at current step */
  MZero( systemVirial.u ); /* initialise total system virial for a new estimation at current step */

  if ( nLipids ) StoreWaterCoorsPrevStep(); // for transbilayer water permeability measurements

  LeapfrogStep( 1 );       

  if ( applyBarostat ) {
    ScaleCoords(); 
    ApplyBoundaryCondS();
    UpdateCellSize(); 
    UnscaleCoords();
  } else ApplyBoundaryCond();

  DO_SITE { 
    VZero( siteForce[ n ] );  // initialize forces
    VZero( siteTorque[ n ] );  // initialize torques
  }

  ComputeNonBondedInteractions( &potEnergy.val, &systemVirial ); 

  if ( nLipids ) { // evaluate intramolecular interactions
    ComputeBondStretchInteractions( &potEnergy.val, &systemVirial );
    ComputeAngleBendInteractions( &potEnergy.val, &systemVirial ); 
    ComputeOrientationRestraints( &potEnergy.val );
  }

  if ( zConstraint ) { /* first part of z-constraint algorithm - assuming there is only 1 solute and it is the [nSites-1] element of the site array */
    fzConst = siteForce[ nSites - 1 ].z; 
    for ( i = 0; i < ( nSites - 1 ); i++ ) siteForce[ i ].z += fzConst / ( nSites - 1 ); /* [J Phys Chem B 2009, 113, 12019, Supporting Information, eq. 5] */
    //      for ( i = 0; i < ( nSites - 1 ); i++ ) siteForce[ i ].z += mass[ site[ i ].type ] * fzConst / totMassExceptConstrSolute; // mass-weighted version
    zConstraintForce.val = - fzConst;  
    PrintZConstraintForce(); // printing to file
    siteForce[ nSites - 1 ].z = 0.; // reset z-force on z-constrained solute
  }

  LeapfrogStep( 2 );   

  if ( nLipids ) { 
    StoreWatCoordsCurrStep();
    CountPossiblePermeatingWaters();
  }

  if ( zConstraint ) { /* second part of z-constraint algorithm */
    vzConst = mass[ site[ nSites - 1 ].type ] * siteVelocity[ nSites -1 ].z / totMassExceptConstrSolute; 
    for ( i = 0; i < ( nSites - 1 ); i++ ) siteVelocity[ i ].z += vzConst; // [J Phys Chem B 2009, 113, 12019, Supporting Information, eq. 4] 
    siteVelocity[ nSites -1 ].z = 0.; // reset z-velocity on z-constrained solute
  } 
  
  if ( removeSystemTranslation ) RemoveSystemTranslation(); 
  if ( removeMonolayersTranslation ) RemoveMonolayersTranslation(); 

  if ( applyThermostat ) ApplyThermostat();

  EvalProps( systemVirial, &systemPressure ); 

  //    CheckPressureTensor( systemPressure );

  if ( applyBarostat ) ApplyBarostat( systemPressure );

  if ( stepCount % stepNebr == 0 ) BuildNeighborList();

  AccumProps( 1 );
  
  if ( stepCount % stepAvg == 0 ) { // compute averages
    // PrintMatrix(boxMatrix);
    AccumProps( 2 );
#ifdef DEBUG
    printf("##step %d: ", stepCount);
    printf("totEn = % 4.0g ", totEnergy.sum ); 
    printf("potEn = % 4.0g ", potEnergy.sum ); 
    printf("kinen = % 4.0g\n", temp.sum * ( kB_IN_BRAHMS_UNITS * DOFs ) / 2.); 
#endif  
    if ( nLipids ) {
      PrintSummary( stdout );
      PrintSummaryOnFile();
      PrintTailOrderParametersOnFile();
      if ( zConstraint ) PrintAvgZConstraintParameters();
      //      PrintSplitTemperatures( stdout );
    } else if ( 0 == nLipids && nWaters ) { // SSD water systems
      PrintSummaryWater( stdout );
      PrintSummaryWaterOnFile();
    } 
    AccumProps( 0 );
  }
}

  
  
  
      
    
  
